import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QTextEdit, QApplication
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar


class MapVisualizerWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.file_paths = [
            'data/quadrant_1.txt',
            'data/quadrant_2.txt',
            'data/quadrant_3.txt',
            'data/quadrant_4.txt'
        ]
        self.init_ui()
        self.load_data()
        self.create_visualization()

    def init_ui(self):
        self.setWindowTitle("3D 环境可视化")
        self.setGeometry(100, 100, 1024, 768)

        main_layout = QVBoxLayout()

        # 创建3D画布
        self.fig = plt.figure()
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.canvas = FigureCanvas(self.fig)
        self.toolbar = NavigationToolbar(self.canvas, self)

        # 创建控制部件
        control_layout = QHBoxLayout()
        self.input_label = QLabel("层/行/列/值（空格分隔）:")
        self.input_field = QLineEdit()
        self.submit_btn = QPushButton("修改障碍物")
        self.submit_btn.clicked.connect(self.handle_edit)

        control_layout.addWidget(self.input_label)
        control_layout.addWidget(self.input_field)
        control_layout.addWidget(self.submit_btn)

        # 信息显示区域
        self.info_text = QTextEdit()
        self.info_text.setReadOnly(True)

        main_layout.addWidget(self.toolbar)
        main_layout.addWidget(self.canvas)
        main_layout.addLayout(control_layout)
        main_layout.addWidget(self.info_text)

        self.setLayout(main_layout)

    def load_data(self):
        """加载四象限数据"""
        self.map_data = []
        self.agent_info = []

        for path in self.file_paths:
            with open(path) as f:
                # 读取地图尺寸
                rows, cols = map(int, f.readline().split())
                # 读取地图数据
                matrix = []
                for _ in range(rows):
                    line = f.readline().strip().split()
                    matrix.append([1 if c == '@' else 0 for c in line])
                # 读取Agent信息
                agent_count = int(f.readline())
                agents = []
                for _ in range(agent_count):
                    agents.append(tuple(map(int, f.readline().split())))

            self.map_data.append(np.array(matrix))
            self.agent_info.append(agents)

    def create_visualization(self):
        """创建初始可视化"""
        self.ax.clear()

        # 绘制障碍物
        for layer, data in enumerate(self.map_data):
            rows, cols = data.shape
            for i in range(rows):
                for j in range(cols):
                    if data[i, j] == 1:
                        self.create_cube(layer, i, j)

        # 设置坐标轴
        self.ax.set_xlabel('X')
        self.ax.set_ylabel('Y')
        self.ax.set_zlabel('层级')
        self.canvas.draw()

    def create_cube(self, layer, row, col):
        """创建立方体"""
        vertices = [
            [col, row, layer],
            [col + 1, row, layer],
            [col + 1, row + 1, layer],
            [col, row + 1, layer],
            [col, row, layer + 1],
            [col + 1, row, layer + 1],
            [col + 1, row + 1, layer + 1],
            [col, row + 1, layer + 1]
        ]
        faces = [
            [vertices[0], vertices[1], vertices[2], vertices[3]],
            [vertices[4], vertices[5], vertices[6], vertices[7]],
            [vertices[0], vertices[1], vertices[5], vertices[4]],
            [vertices[1], vertices[2], vertices[6], vertices[5]],
            [vertices[2], vertices[3], vertices[7], vertices[6]],
            [vertices[3], vertices[0], vertices[4], vertices[7]]
        ]
        poly = Poly3DCollection(faces, facecolor='#808080', edgecolor='k', alpha=0.8)
        self.ax.add_collection3d(poly)

    def handle_edit(self):
        """处理障碍物编辑"""
        try:
            parts = self.input_field.text().split()
            if len(parts) != 4:
                raise ValueError("需要4个参数")

            layer = int(parts[0])
            row = int(parts[1])
            col = int(parts[2])
            value = int(parts[3])

            if not (0 <= layer < 4):
                raise ValueError("层级需在0-3之间")
            if value not in (0, 1):
                raise ValueError("值只能是0或1")

            # 更新数据
            self.map_data[layer][row, col] = value
            self.update_visualization()
            self.save_changes(layer)

        except Exception as e:
            self.info_text.append(f"错误: {str(e)}")

    def update_visualization(self):
        """更新可视化"""
        self.create_visualization()

    def save_changes(self, layer):
        """保存修改到文件"""
        path = self.file_paths[layer]
        with open(path, 'w') as f:
            # 写入尺寸
            rows, cols = self.map_data[layer].shape
            f.write(f"{rows} {cols}\n")

            # 写入地图数据
            for row in self.map_data[layer]:
                line = ' '.join(['@' if x == 1 else '.' for x in row])
                f.write(line + '\n')

            # 写入Agent信息
            f.write(f"{len(self.agent_info[layer])}\n")
            for agent in self.agent_info[layer]:
                f.write(f"{' '.join(map(str, agent))}\n")

        self.info_text.append(f"层级 {layer} 已保存")